from modelscope.pipelines import pipeline
from modelscope.utils.constant import Tasks
from configs.model_config import MODEL_PATH, FAQ_THRESHOLD, FAQ_MODEL
import time
from typing import NamedTuple

# FAQ模型
FAQ = pipeline(Tasks.faq_question_answering, MODEL_PATH["faq_model"][FAQ_MODEL], device="cpu")


# 定义一个命名元组来表示函数的返回类型
class FAQResult(NamedTuple):
    spend_time: float
    label: str
    probability: float
    answer: str


def faq(**kwargs) -> FAQResult:
    # FAQ开始时间
    FAQ_start_time = time.time()

    # 从kwargs中获取必要的参数
    query = kwargs.get('query', None)
    QA_datas = kwargs.get('QA_datas', None)
    KV_datas = kwargs.get('KV_datas', None)
    FAQ_MODEL = kwargs.get('FAQ_MODEL', None)

    # 确保所有必要的参数都已提供
    if query is None or QA_datas is None or FAQ_MODEL is None:
        raise ValueError("必要的参数query、QA_datas和FAQ_MODEL必须提供。")

    # FAQ模型计算
    output = FAQ({"query_set": [query], "support_set": QA_datas})

    # FAQ的概率
    FAQ_probability = output['output'][0][0]['score']

    # FAQ的标签
    FAQ_label = output['output'][0][0]['label']

    # FAQ结束时间
    FAQ_end_time = time.time()

    # FAQ耗时
    FAQ_spend_time = FAQ_end_time - FAQ_start_time
    # FAQ答案
    FAQ_answer = KV_datas[FAQ_label]

    # 使用命名元组作为返回类型
    return FAQResult(spend_time=FAQ_spend_time, label=FAQ_label, probability=FAQ_probability, answer=FAQ_answer)
